import numpy as np
from dataset.mnist import load_mnist
from two_layer_net import TwoLayerNet
import matplotlib.pyplot as plt

# 加载数据集
(x_train, t_train), (x_test, t_test) = load_mnist(normalize=True, one_hot_label=True)

if __name__ == '__main__':
    train_loss_list = []
    # 梯度法更新次数 10000>20mins
    iters_num = 10000
    train_size = x_train.shape[0]
    batch_size = 100
    learn_rate = 0.1

    network = TwoLayerNet(input_size=784, hidden_size=50, output_size=10)

    # 梯度更新 iters_num 更新次数
    for i in range(iters_num):

        # 获取mini_batch
        batch_mask = np.random.choice(train_size, batch_size)
        x_batch = x_train[batch_mask]
        t_batch = t_train[batch_mask]

        # TODO 计算梯度
        # 非常的慢 速度令人发指
        # grad = network.numerical_gradient(x_batch, t_batch)
        grad = network.gradient(x_batch, t_batch) # 高速版（误差反向传播算法）

        # TODO 更新参数
        for key in ('W1', 'b1', 'W2', 'b2'):
            network.params[key] -= learn_rate * grad[key]

        # TODO 记录学习过程loss函数的变化
        loss = network.loss(x_batch, t_batch)
        train_loss_list.append(loss)

    # 可视化结果损失函数变化结果
    print(train_loss_list)
    plt.figure(figsize=(10,5))
    plt.plot(train_loss_list,label = 'Training Loss')
    plt.title('Training Loss over Time')
    # poch是神经网络训练中的一个周期，
    # 表示整个训练数据集被模型遍历一次的过程。
    plt.xlabel('Iterations / Epoch')
    plt.ylabel('Loss')
    # 显示图例
    plt.legend()
    # 显示网格
    plt.grid(True)
    plt.savefig("train_loss_plot.png")
    plt.show()


